魔性!Python生成全网爆火的“蚂蚁呀嘿”

您所在的位置:网站首页 read mp4 魔性!Python生成全网爆火的“蚂蚁呀嘿”

魔性!Python生成全网爆火的“蚂蚁呀嘿”

#魔性!Python生成全网爆火的“蚂蚁呀嘿”| 来源: 网络整理| 查看: 265

魔性!Python生成全网爆火的“蚂蚁呀嘿”

前段时间火爆的“蚂蚁呀嘿”,将一个人的说话动作和表情迁移到另一张静态图中,让静态图中的人脸做出指定动作表情,主要基于FOMM(First Order Motion model)技术。这已经是2年前的技术了,在一些场景中生成的效果并不理想。近期,清华大学团队在CVPR2022发布最新表情动作迁移论文Thin-Plate Spline Motion Model for Image Animation。本文不具体讲论文原理,而是直接将其开源的模型down下来用。效果如下:

第1张图是静态照片,第二张是gif驱动动画,第三张是生成的结果。

本文目的:将开源模型打包封装成一个单独接口,读者只需简单传入一张图片和一个动画(gif或mp4),即可生成表情迁移动画(mp4)。

0 环境准备

读者需要安装好pytorch环境,可前往https://pytorch.org/get-started/locally/根据实际的硬件环境,选择GPU或cpu版本。

安装imageio-ffmpeg库,用于读取mp4文件。

读者可以直接跳到最后,获取源码,下载源码资源包后,将文件夹中的对应图片替换成自己的图片,即可一键生成“蚂蚁呀嘿”。

1 模型封装

将模型导出为pt后,再创建Model类,将接口封装到infer函数, 具体代码如下:

class Model(): def __init__(self, kp="models/kp.pt", aio="models/aio.pt", device=torch.device('cpu')): self.device = device self.kp = torch.jit.load(kp, map_location=device).eval() self.aio = torch.jit.load(aio, map_location=device).eval() def relative_kp(self, kp_source, kp_driving, kp_driving_initial): source_area = ConvexHull(kp_source[0].data.cpu().numpy()).volume driving_area = ConvexHull( kp_driving_initial[0].data.cpu().numpy()).volume adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area) kp_new = kp_driving kp_value_diff = (kp_driving - kp_driving_initial) kp_value_diff *= adapt_movement_scale kp_new = kp_value_diff + kp_source return kp_new def get_kp(self, src): src = np.expand_dims(src, 0).transpose(0, 3, 1, 2) src = torch.from_numpy(src).float().to(self.device) return self.kp(src) def infer(self, src, driving, src_kp, init_kp): src = np.expand_dims(src, 0).transpose(0, 3, 1, 2) src = torch.from_numpy(src).float().to(self.device) driving = np.expand_dims(driving, 0).transpose(0, 3, 1, 2) driving = torch.from_numpy(driving).float().to(self.device) kp_driving = self.kp(driving) kp_norm = self.relative_kp(kp_source=src_kp, kp_driving=kp_driving, kp_driving_initial=init_kp) with torch.no_grad(): out = self.aio(src, src_kp, kp_norm) out = out[0].cpu().numpy() out = out.transpose(1, 2, 0) return out

其中,get_kp函数用于获取脸部关键点数据。infer函数中,src表示静态图,driving表示动态图中的某一帧,src_kp表示静态图的关键点,init_kp表示动态图中的第一帧的关键点。

2 调用模型

整个调用流程可以拆分为4步:创建模型对象、读取动图的每一帧、调用模型、生成帧导出mp4。

2.1 创建模型对象

前面定义好了Model对象,需要根据GPU和CPU环境,由读者指定使用具体的pytorch版本,具体代码如下所示。

def create_model(use_gpu): if use_gpu: device = torch.device('cuda') else: device = torch.device('cpu') model = Model(device=device) return model

上面代码中,use_gpu是个boolean类型,用于判断是否使用GPU版本,读者根据自己的时间情况设置。

2.2 读取动图的每一帧

调用imageio-ffmpeg库,读取mp4或gif文件中的每一帧。具体代码如下所示,函数返回列表,列表内容为视频帧:

def read_mp4_or_gif(path): reader = imageio.get_reader(path) if path.lower().endswith('.mp4'): fps = reader.get_meta_data().get('fps') elif path.lower().endswith('.gif'): fps = 1000 / Image.open(path).info['duration'] driving_video = [] try: for im in reader: im = resize(im, (256, 256))[..., :3] driving_video.append(im) except RuntimeError: pass reader.close() return driving_video, fps

因为模型的约束,这里将每一帧resize到256*256。

2.3 调用模型

模型调用非常简单,只需读取静态图和动态图的每一帧,并对静态图和动态图第一帧调用Model类的get_kp函数获取关键点即可。遍历动态图的每一帧,将动图帧、静态图、静态图关键点、动图第一帧关键点一起传给Model的infer函数即可得到生成的帧。具体代码如下所示。

def run(use_gpu, src_path, driving_path): src = imageio.imread(src_path) src = resize(src, (256, 256))[..., :3] driving_video, fps = read_mp4_or_gif(driving_path) model = create_model(use_gpu) src_kp = model.get_kp(src) init_kp = model.get_kp(driving_video[0]) outs = [] for driving in driving_video: out = model.infer(src, driving, src_kp, init_kp) out = img_as_ubyte(out) outs.append(out) return outs, fps 2.4 生成帧导出mp4

这里继续调用imageio-ffmpeg库,将视频帧组装成mp4文件,代码如下:

def write_mp4(out_path, frames, fps): imageio.mimsave(out_path, frames, fps=fps) 2.5 pipeline调用

整个调用流水线如下:

src_path = 'assets/source.png' driving_path = 'assets/driving2.gif' frames, fps = run(True, src_path, driving_path) write_mp4("out.mp4", frames, fps) 3 源码获取关注公众号:Python学习实战公众号聊天界面回复:表情迁移,获取完整源码。


【本文地址】


今日新闻


推荐新闻


CopyRight 2018-2019 办公设备维修网 版权所有 豫ICP备15022753号-3